# -*- coding:utf-8 -*-
import tensorflow as tf
import os
from tensorflow.examples.tutorials.mnist import input_data
#加载mnist_inference.py中定义的常量和前向传播的函数
import mnist_inference

#配置神经网络的参数（超参数，而不是网络结构的参数）
BATCH_SIZE=100
LEARNING_RATE_BASE=0.8  #初始学习率
LEARNING_RATE_DECAY=0.99      #用于指数衰减学习率的衰减系数
REGULARAZTION_RATE=0.0001     #正则化损失中的lambda
TRAINING_STEPS=30000          #训练步数
MOVING_AVERAGE_DECAY=0.99     #滑动平均的衰减率
#模型保存的路径
MODEL_SAVE_PATH='/path/to/model'
MODEL_NAME='model.ckpt'

def train(mnist):
    #定义输入输出placeholder
    x=tf.placeholder(
        tf.float32,[None,mnist_inference.INPUT_NODE],name='x-input'
    )
    y_=tf.placeholder(
        tf.float32,[None,mnist_inference.OUTPUT_NODE],name='y-input'
    )

    #正则化器，之后可以直接regularizer(var)得到变量的正则化值
    regularizer=tf.contrib.layers.l2_rugularizer(REGULARAZTION_RATE)

    y=mnist_inference.inference(x,regularizer)
    global_step=tf.Variable(0,trainable=False)

    #和5.2.1节样例中类似的定义损失函数，学习率，滑动平均操作以及训练过程
    variable_average=tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY,global_step)
    variable_average_op=variable_average.apply(tf.trainable_variables())  #定义滑动平均操作

    cross_entropy=tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y,labels=tf.argmax(y_,1))  #交叉熵
    cross_entropy_mean=tf.reduce_mean(cross_entropy)       #当前batch中所有样例的交叉熵平均值
    loss=cross_entropy_mean+tf.add_n(tf.get_collection('losses'))  #tf.add_n()是将输入按照元素相加

    learning_rate=tf.train.exponential_decay(
        LEARNING_RATE_BASE,global_step,mnist.train.num_example/BATCH_SIZE,LEARNING_RATE_DECAY
    )
    #在minimize函数中传入global_step将自动更新global_step参数，从而使得学习率也得到对应更新。
    train_step=tf.train.GradientDescentOptimizer(learning_rate).minimize(loss,global_step=global_step)

    #组合操作使得一次完成多个操作，下列两行代码等价于：train_op=tf.group(train_step,variable_average_op)
    with tf.control_dependencies([train_step,variable_average_op]):
        train_op=tf.no_op(name='train')


    #初始化TensorFLow持久化类
    saver=tf.train.Saver()
    with tf.Session() as sess:
        tf.global_variables_initializer().run()

        #在训练过程中不再预测模型在验证数据上的表现，验证和测试的过程会有一个独立的程序mnist_eval.py来完成
        for i in range(TRAINING_STEPS):
            xs,ys=mnist.train.next_batch(BATCH_SIZE)
            _,loss_value,step=sess.run([train_op,loss,global_step],feed_dict={x:xs,y_:ys})
            #每1000轮保存一次模型
            if i%1000==0:
                #输出当前的训练情况。这里只输出了模型在当前batch上的损失函数大小，
                # 通过损失函数的大小可以大概了解训练的情况，在验证数据集上的正确率信息会有一个单独的程序来生成。
                print("After %d training step(s),loss on training batch is %g."%(step,loss_value))

                #保存当前的模型。注意这里给出了global_step参数，这样可以让每个被保存的模型文件名末尾加上训练的轮数。
                # 比如'model.ckpt-1000'表示训练1000轮后得到的模型。
                saver.save(sess,os.path.join(MODEL_SAVE_PATH,MODEL_NAME),global_step=global_step)

def main(argv=None):
    mnist=input_data.read_data_sets('/path/to/mnist_data',one_hot=True)
    train(mnist)

#Tensorflow提供的一个主程序入口，tf.app.run()会调用上面定义的main函数
if __name__=='__main__':
    tf.app.run()